"""
@Filename       : matrix.py
@Create Time    : 2021/11/16 21:01
@Author         : Rylynn
@Description    : 

"""
import time
from timeit import timeit

import scipy.linalg as la
import scipy.sparse.linalg as sla
import numpy as np



def svd(m, k):
    return sla.svds(m, k)


def bk_svd(tensor, k=6, block_size=None, num_iter=3):
    """
        Block Krylov SVD
        Parameters
        ----------
        tensor : array (M, N) array_like
        k : int
        Returns
        -------
        u : (M, k) array_like
        s : (k,) array_like
        vt : (k, N) array_like
        """
    if block_size is None:
        block_size = k

    k = min(k, min(tensor.shape))
    u = np.zeros((1, tensor.shape[1]))

    l = np.ones((tensor.shape[0], 1))

    K = np.zeros((tensor.shape[1], block_size * num_iter))
    block = np.random.randn(tensor.shape[1], block_size)

    block, _ = la.qr(block, mode='economic')

    T = np.zeros((tensor.shape[1], block_size))

    for i in range(num_iter):
        T = tensor @ block - l * (u @ block)
        block = tensor.T @ T - (u.T * (l.T @ T))
        block, _ = la.qr(block, mode='economic')
        K[:, i * block_size: (i + 1) * block_size] = block
    Q, _ = la.qr(K, mode='economic')

    # Rayleigh-Ritz
    T = tensor @ Q - l @ (u @ Q)

    Ut, St, Vt = np.linalg.svd(T, full_matrices=False)
    U = Ut[:, :k]
    S = St[:k]
    V = Q @ Vt.T
    return U, S, V[:, :k].T


if __name__ == '__main__':
    a = np.random.random((10000, 5000))
    st = time.time()
    U, S, V = bk_svd(a, k=64)
    et = time.time()
    print(et - st)
    st = time.time()
    U_, S_, V_ = svd(a)
    et = time.time()
    print(et - st)
    print(U)
    print(S)
    print(S_)
    print(V)

